import os
import numpy as np
import re

def save_flow_part(flow_part, row, col, output_dir):
    filename = f"tr{row + 1}-tc{col + 1}.npy"
    filepath = os.path.join(output_dir, filename)
    np.save(filepath, flow_part)
    print(f"Saved {filename} to {filepath}")

def merge_and_save_flows(flow_dir, output_dir, tile_size=256, small_tile_size=1024):
    # 获取所有光流场文件
    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]

    # 解析所有光流场文件的坐标
    coordinates = []
    for flow_file in flow_files:
        match = re.search(r'tile_\((\d+)_(\d+)\)\.npy', flow_file)
        if not match:
            raise ValueError(f"文件名格式不正确: {flow_file}")
        row, col = map(int, match.groups())
        coordinates.append((row, col))

    # 确定大光流场的尺寸
    max_row = max(row for row, col in coordinates)
    max_col = max(col for row, col in coordinates)

    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 初始化存储1024x1024块的变量
    for block_row in range(0, max_row + 1, small_tile_size // tile_size):
        for block_col in range(0, max_col + 1, small_tile_size // tile_size):
            flow_block = np.zeros((small_tile_size, small_tile_size, 2), dtype=np.float32)
            for i in range(small_tile_size // tile_size):
                for j in range(small_tile_size // tile_size):
                    row = block_row + i
                    col = block_col + j
                    if (row, col) in coordinates:
                        flow_file = f"tile_({row}_{col}).npy"
                        flow_path = os.path.join(flow_dir, flow_file)
                        flow = np.load(flow_path)
                        start_row = i * tile_size
                        start_col = j * tile_size
                        flow_block[start_row:start_row + tile_size, start_col:start_col + tile_size] = flow
            save_flow_part(flow_block, (int)(block_row/4), (int)(block_col/4), output_dir)

# 示例用法
flow_dir = r'C:\\11867_32nm_fine_align_256_dvf_repair(3)'  # 替换为实际路径
output_dir = r'E:\wafer52\32nm_flow_fine_align(3)'  # 替换为实际输出路径

merge_and_save_flows(flow_dir, output_dir, tile_size=256, small_tile_size=1024)

print(f"Processed and saved flow parts to {output_dir}")
